import numpy as np
import random as rd
import drawNet as dn
def WS(n, k, p):#WS小世界
    G = np.empty((n, n))
    lst = []
    for i in range(0, n):
        G[i][i] = -1
        for j in range(i + 1, n):
            G[i][j] = G[j][i] = -1

    for i in range(0, n):
        for j in range(0, int(k / 2)):
            nxt = (i + j + 1) % n 
            G[i][nxt] = G[nxt][i] = 1
            lst.append([nxt, i])
            lst.append([i, nxt])
    
    while(len(lst) > 0):
        k = rd.choice(lst)
        r = rd.random()
        if(r < p):
            cp_lst = []
            for i in range(0, n):
                if(i != k[0] and G[k[0]][i] == -1):
                    cp_lst.append(i)
            if(len(cp_lst) != 0):
                rdp = rd.choice(cp_lst)
                G[k[0]][k[1]] = G[k[1]][k[0]] = -1
                G[k[0]][rdp] = G[rdp][k[0]] = 1
        
        lst.remove(k)
        lst.remove([k[1], k[0]])
    return G
dn.drawGraph(WS(10, 4, 0))
dn.drawGraph(WS(20, 4, 0.3))